import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

# data sets
df = pd.read_csv("HeightWeight.csv", delimiter=',')
x = df.iloc[:, 1:3].values
print("x: ", x.shape)

m, n = x.shape

# random center
K = 2
from random import sample
rand_indices = sample(range(0, m),K)
center = np.array([x[i] for i in rand_indices])

# data visualization
plt.scatter(x[:, 0], x[:,1])
plt.plot(center[:, 0], center[:, 1], 'rx')
plt.show()

# square ditance
def ditanceSquare(point1, point2):
    dist = np.sum(np.square(point1 - point2))
    return dist

# find closest center
def findClosestCenter(X, center):
    index = np.zeros([m, 1])
    for i in range(m):
        mindist, index_temp = 100000, 0
        for k in range(K):
            dist = ditanceSquare(center[k], X[i])
            if dist < mindist:
                mindist = dist
                index_temp = k
        index[i] = index_temp
    return index

# compute center
def computeCenter(X, index):
    clustX = []
    for k in range(K):
        clustX.append(np.array([X[i] for i in range(m) if index[i] == k]))
    center = np.array([np.mean(thisX, axis=0) for thisX in clustX])
    return clustX, center

# hyper parameter
iters = 10
def Kmeans(X, center):
    center_his = []
    for step in range(iters):
        center_his.append(center)
        index = findClosestCenter(X, center)
        _, center = computeCenter(X, index)
    return index, center_his

# main program
index, center_his = Kmeans(x, center)
print("center:", center_his[-1])

# data plot
def plotData(X, index, center_his):
    clustX, _ = computeCenter(X, index)
    for k in range(K):
        plt.scatter(clustX[k][:, 0], clustX[k][:, 1])
        plt.plot([center[k, 0] for center in center_his],
                 [center[k, 1] for center in center_his], 'rx--')

plotData(x, index, center_his)
plt.show()
